from langchain_core.language_models import BaseChatModel

from domain.model.platform.base import BasePlatformChatModel


class ChatModelManager:

    def __init__(self):
        # 平台模型映射
        self._model_ins_map: dict[str, BasePlatformChatModel] = {}

    def register_platform_model(self, platform: str, platform_model_instance: BasePlatformChatModel):
        """注册平台模型"""
        if self._model_ins_map.get(platform):
            raise Exception(f"{platform} platform is already registered")
        self._model_ins_map[platform] = platform_model_instance

    def get_model_instance_by_name(self, model_name: str) -> BaseChatModel:
        """通过模型名称获取模型实例"""
        for platform_model in self._model_ins_map.values():
            if model_name in platform_model.get_support_model_list():
                return platform_model.get_model_ins_by_name(model_name)
        raise Exception(f'cannot find model named {model_name}')
